#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
昇腾处理器工具函数
提供昇腾NPU初始化、资源管理等功能
"""

import os
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# 检查是否有昇腾环境
try:
    import acl
    import torch_npu
    HAS_ASCEND = True
except ImportError:
    HAS_ASCEND = False
    logger.warning("未检测到昇腾SDK，将使用CPU/GPU模式运行")


def init_ascend(device_id: int = 0) -> bool:
    """
    初始化昇腾环境
    
    Args:
        device_id: 设备ID
        
    Returns:
        初始化是否成功
    """
    if not HAS_ASCEND:
        logger.warning("未检测到昇腾SDK，无法初始化昇腾环境")
        return False
    
    try:
        # 设置环境变量
        os.environ['ASCEND_DEVICE_ID'] = str(device_id)
        
        # 初始化ACL
        ret = acl.init()
        if ret != 0:
            logger.error(f"ACL初始化失败，错误码: {ret}")
            return False
        
        # 设置当前设备
        ret = acl.rt.set_device(device_id)
        if ret != 0:
            logger.error(f"设置设备失败，错误码: {ret}")
            acl.finalize()
            return False
        
        # 创建上下文
        context = acl.rt.create_context(device_id)
        
        # 初始化PyTorch NPU
        torch_npu.npu.set_device(f"npu:{device_id}")
        
        logger.info(f"昇腾环境初始化成功，设备ID: {device_id}")
        return True
    
    except Exception as e:
        logger.error(f"初始化昇腾环境时出错: {str(e)}")
        return False


def release_ascend(device_id: int = 0) -> bool:
    """
    释放昇腾资源
    
    Args:
        device_id: 设备ID
        
    Returns:
        释放是否成功
    """
    if not HAS_ASCEND:
        return False
    
    try:
        # 销毁上下文
        ret = acl.rt.destroy_context(acl.rt.get_context())
        if ret != 0:
            logger.error(f"销毁上下文失败，错误码: {ret}")
            return False
        
        # 重置设备
        ret = acl.rt.reset_device(device_id)
        if ret != 0:
            logger.error(f"重置设备失败，错误码: {ret}")
            return False
        
        # 结束ACL
        ret = acl.finalize()
        if ret != 0:
            logger.error(f"ACL结束失败，错误码: {ret}")
            return False
        
        logger.info(f"昇腾资源释放成功，设备ID: {device_id}")
        return True
    
    except Exception as e:
        logger.error(f"释放昇腾资源时出错: {str(e)}")
        return False


def get_npu_memory_usage(device_id: int = 0) -> Optional[dict]:
    """
    获取NPU内存使用情况
    
    Args:
        device_id: 设备ID
        
    Returns:
        包含内存使用情况的字典，失败则返回None
    """
    if not HAS_ASCEND:
        return None
    
    try:
        # 获取设备内存信息
        free_mem, total_mem = torch_npu.npu.get_device_properties(device_id).total_memory, torch_npu.npu.max_memory_allocated(device_id)
        used_mem = total_mem - free_mem
        
        return {
            "total": total_mem,
            "used": used_mem,
            "free": free_mem,
            "unit": "Byte"
        }
    
    except Exception as e:
        logger.error(f"获取NPU内存使用情况时出错: {str(e)}")
        return None


def optimize_model_for_ascend(model) -> bool:
    """
    优化模型以适配昇腾处理器
    
    Args:
        model: PyTorch模型
        
    Returns:
        优化是否成功
    """
    if not HAS_ASCEND or model is None:
        return False
    
    try:
        # 启用JIT编译
        torch_npu.npu.set_compile_mode(jit_compile=True)
        
        # 如果模型支持半精度，转为FP16
        if hasattr(model, "half"):
            model.half()
        
        # 设置动态输入
        torch_npu.npu.set_aoe(mode="dynamic_dims")
        
        # 优化内存分配
        torch_npu.npu.set_memory_fraction(0.8)
        
        return True
    
    except Exception as e:
        logger.error(f"优化模型时出错: {str(e)}")
        return False


def get_available_npu_devices() -> list:
    """
    获取可用的昇腾设备列表
    
    Returns:
        可用设备ID列表
    """
    if not HAS_ASCEND:
        return []
    
    try:
        # 获取设备数量
        device_count = torch_npu.npu.device_count()
        return list(range(device_count))
    
    except Exception as e:
        logger.error(f"获取可用NPU设备时出错: {str(e)}")
        return [] 